# -*- coding: utf-8 -*-
"""
@Time    : 2024/6/28 11:17 
@Author  : ZhangShenao 
@File    : 2.使用RunnableParallel模拟RAG.py 
@Desc    : 使用RunnableParallel模拟RAG
RunnableParallel可以用于操作Runnable的输出,以产生符合下一个Runnable组件的数据
"""
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableParallel, RunnablePassthrough


def retrieval(query: str) -> str:
    """
    模拟检索操作,根据用户的输入搜索结果并返回
    :param query: 用户输入
    :return: 检索结果
    """
    print(f'query = {query}, 执行检索操作.....')
    return '我叫zsa，是一名资深AI应用开发工程师'


# 编排Prompt
prompt = ChatPromptTemplate.from_template('''请参考相关的上下文，并根据用户的提问回答问题：
<context>
{context}
</context>

用户提问: {question}
''')

# 创建智普AI LLM
llm = ChatZhipuAI(model="GLM-4-Air")

# 创建OutputParser
parser = StrOutputParser()

# 构建Chain,并使用RunnableParallel操作组件的输出结果
chain = RunnableParallel(
    context=retrieval,
    question=RunnablePassthrough(),  # RunnablePassthrough可以透传上游的输入,简化Chain的调用
) | prompt | llm | parser

# 执行Chain
result = chain.invoke('我是谁？')

# 输出结果
print(result)  # 根据您提供的上下文，您是zsa，一名资深AI应用开发工程师。
